#include "gpu.h"
#include <cuda_runtime.h>
#include <iostream>
#include <cstring>
#include <chrono> // 添加计时功能

// 全局变量声明
int total_guesses = 0;

// CUDA内核函数：处理单段PT的猜测生成
__global__ void processStringKernel(const char* prefix, int prefix_len, 
                                  const char* values_data, int* lengths, int* offsets,
                                  int num, char* output, int* output_offsets) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (idx < num) {
        int val_len = lengths[idx];
        int val_offset = offsets[idx];
        int result_len = prefix_len + val_len;
        int out_offset = output_offsets[idx];
        
        // 复制前缀字符串
        for (int i = 0; i < prefix_len; i++) {
            output[out_offset + i] = prefix[i];
        }
        
        // 复制值字符串
        for (int i = 0; i < val_len; i++) {
            output[out_offset + prefix_len + i] = values_data[val_offset + i];
        }
        
        // 添加终止符
        output[out_offset + result_len] = '\0';
    }
}

void GPU(std::vector<std::string> &val, int num, std::vector<std::string> &guesses, const std::string &guess) {
    if (num <= 0 || val.empty()) return;

    // 性能分析变量
    auto start_total = std::chrono::high_resolution_clock::now();
    auto start_prep = std::chrono::high_resolution_clock::now();
    auto end_prep = std::chrono::high_resolution_clock::now();
    auto start_mem_alloc = std::chrono::high_resolution_clock::now();
    auto end_mem_alloc = std::chrono::high_resolution_clock::now();
    auto start_h2d = std::chrono::high_resolution_clock::now();
    auto end_h2d = std::chrono::high_resolution_clock::now();
    auto start_kernel = std::chrono::high_resolution_clock::now();
    auto end_kernel = std::chrono::high_resolution_clock::now();
    auto start_d2h = std::chrono::high_resolution_clock::now();
    auto end_d2h = std::chrono::high_resolution_clock::now();
    auto end_total = std::chrono::high_resolution_clock::now();

    // 开始预处理计时
    start_prep = std::chrono::high_resolution_clock::now();

    // 准备主机端数据
    const char* prefix = guess.c_str();
    int prefix_len = guess.length();
    
    // 为扁平化的字符串数据预先计算总长度
    int values_total_len = 0;
    int* h_lengths = new int[num];
    int* h_offsets = new int[num];
    int* h_output_offsets = new int[num];
    
    int output_total_len = 0;
    
    for (int i = 0; i < num; i++) {
        h_lengths[i] = val[i].length();
        h_offsets[i] = values_total_len;
        h_output_offsets[i] = output_total_len;
        
        values_total_len += h_lengths[i];
        output_total_len += prefix_len + h_lengths[i] + 1; // +1 for null terminator
    }
    
    // 创建扁平化的字符串数据
    char* h_values_data = new char[values_total_len];
    
    // 复制所有字符串数据到扁平化的数组中
    for (int i = 0; i < num; i++) {
        memcpy(h_values_data + h_offsets[i], val[i].c_str(), h_lengths[i]);
    }
    
    // 结束预处理计时
    end_prep = std::chrono::high_resolution_clock::now();
    
    // 开始内存分配计时
    start_mem_alloc = std::chrono::high_resolution_clock::now();
    
    // 一次性分配所有设备内存
    char* d_prefix;
    char* d_values_data;
    int* d_lengths;
    int* d_offsets;
    int* d_output_offsets;
    char* d_output;
    
    cudaMalloc(&d_prefix, prefix_len + 1);
    cudaMalloc(&d_values_data, values_total_len);
    cudaMalloc(&d_lengths, num * sizeof(int));
    cudaMalloc(&d_offsets, num * sizeof(int));
    cudaMalloc(&d_output_offsets, num * sizeof(int));
    cudaMalloc(&d_output, output_total_len * sizeof(char));
    
    // 结束内存分配计时
    end_mem_alloc = std::chrono::high_resolution_clock::now();
    
    // 开始主机到设备数据传输计时
    start_h2d = std::chrono::high_resolution_clock::now();
    
    // 一次性复制所有数据到设备
    cudaMemcpy(d_prefix, prefix, prefix_len + 1, cudaMemcpyHostToDevice);
    cudaMemcpy(d_values_data, h_values_data, values_total_len, cudaMemcpyHostToDevice);
    cudaMemcpy(d_lengths, h_lengths, num * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(d_offsets, h_offsets, num * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(d_output_offsets, h_output_offsets, num * sizeof(int), cudaMemcpyHostToDevice);
    
    // 结束主机到设备数据传输计时
    end_h2d = std::chrono::high_resolution_clock::now();
    
    // 开始内核执行计时
    start_kernel = std::chrono::high_resolution_clock::now();
    
    // 启动内核
    int blockSize = 256;
    int gridSize = (num + blockSize - 1) / blockSize;
    processStringKernel<<<gridSize, blockSize>>>(d_prefix, prefix_len, d_values_data, 
                                             d_lengths, d_offsets, num, d_output, 
                                             d_output_offsets);
    
    // 同步设备
    cudaDeviceSynchronize();
    
    // 结束内核执行计时
    end_kernel = std::chrono::high_resolution_clock::now();
    
    // 开始设备到主机数据传输计时
    start_d2h = std::chrono::high_resolution_clock::now();
    
    // 将结果复制回主机
    char* h_output = new char[output_total_len];
    cudaMemcpy(h_output, d_output, output_total_len, cudaMemcpyDeviceToHost);
    
    // 结束设备到主机数据传输计时
    end_d2h = std::chrono::high_resolution_clock::now();
    
    // 构建结果字符串
    for (int i = 0; i < num; i++) {
        guesses.push_back(std::string(h_output + h_output_offsets[i]));
    }
    
    // 更新全局计数
    total_guesses += num;
    
    // 释放内存
    delete[] h_values_data;
    delete[] h_lengths;
    delete[] h_offsets;
    delete[] h_output_offsets;
    delete[] h_output;
    
    cudaFree(d_prefix);
    cudaFree(d_values_data);
    cudaFree(d_lengths);
    cudaFree(d_offsets);
    cudaFree(d_output_offsets);
    cudaFree(d_output);
    
    // 结束总时间计时
    end_total = std::chrono::high_resolution_clock::now();
    
    // 计算各阶段耗时（毫秒）
    auto prep_time = std::chrono::duration_cast<std::chrono::milliseconds>(end_prep - start_prep).count();
    auto mem_alloc_time = std::chrono::duration_cast<std::chrono::milliseconds>(end_mem_alloc - start_mem_alloc).count();
    auto h2d_time = std::chrono::duration_cast<std::chrono::milliseconds>(end_h2d - start_h2d).count();
    auto kernel_time = std::chrono::duration_cast<std::chrono::milliseconds>(end_kernel - start_kernel).count();
    auto d2h_time = std::chrono::duration_cast<std::chrono::milliseconds>(end_d2h - start_d2h).count();
    auto total_time = std::chrono::duration_cast<std::chrono::milliseconds>(end_total - start_total).count();
    
    // 打印性能分析结果
    static int call_count = 0;
    if (++call_count % 500 == 0) { // 每500次调用打印一次，避免输出过多
        std::cout << "\n==== GPU Performance Analysis (call #" << call_count << ") ====\n";
        std::cout << "Number of strings: " << num << "\n";
        std::cout << "Preprocessing: " << prep_time << " ms (" << (prep_time*100.0/total_time) << "%)\n";
        std::cout << "Memory allocation: " << mem_alloc_time << " ms (" << (mem_alloc_time*100.0/total_time) << "%)\n";
        std::cout << "Host to device transfer: " << h2d_time << " ms (" << (h2d_time*100.0/total_time) << "%)\n";
        std::cout << "Kernel execution: " << kernel_time << " ms (" << (kernel_time*100.0/total_time) << "%)\n";
        std::cout << "Device to host transfer: " << d2h_time << " ms (" << (d2h_time*100.0/total_time) << "%)\n";
        std::cout << "Total GPU function time: " << total_time << " ms\n";
        std::cout << "======================================\n";
    }
}